package com.demo.树;

import lombok.Getter;
import lombok.Setter;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;

/**
 * 字典树 或者说前缀树
 * @author cl
 */
public class Trie {

    private TrieNode head = null;

    @Getter
    private List<String> data = new ArrayList<>();

    @Getter
    private ArrayList<TrieNode> nodes = new ArrayList<>();

    public Trie(String... strings) {
        if (strings == null || strings.length == 0) {
            throw new IllegalArgumentException("参数必须要有效!");
        }
        createTrie(strings);
    }

    /** 创建一个字典树 */
    private void createTrie(String... strings) {
        this.head = new TrieNode();
        this.head.word = "";
        nodes.add(head);
        // 循环构建
        for(String str : strings) {
            addWord(str);
        }
    }

    /** 添加一个单词到字典树中 */
    public void addWord(String word) {
        // null 不处理
        if (word == null) {
            return;
        }
        // 有字符串
        this.head.pass++;
        TrieNode current = this.head;
        for (char c : word.toCharArray()) {
            // 添加一个字符到trie上去
            TrieNode tmp = current.edge[getAscallLetterIndex(c)];
            if (tmp == null) {
                tmp = new TrieNode();
                tmp.word = current.word + c;
                current.edge[getAscallLetterIndex(c)] = tmp;
                nodes.add(tmp);
            }
            tmp.pass++;
            current = tmp;
        }
        current.end++;
        this.data.add(word);
    }

    /** 从字典树中删除一个单词 */
    public void delWord(String word) {
        // null 不处理 word不存在不处理
        if (word == null || !wordExists(word)) {
            return;
        }
        this.data.remove(word);
        // 有字符串
        this.head.pass--;
        if (this.head.pass == 0) {
            this.head = null;
            return;
        }
        TrieNode current = this.head;
        for (char c : word.toCharArray()) {
            // 添加一个字符到trie上去
            TrieNode tmp = current.edge[getAscallLetterIndex(c)];
            if (--tmp.pass == 0) {
                current.edge[getAscallLetterIndex(c)] = null;
                // 做删除
                while (tmp != null) {
                    nodes.remove(tmp);
                    for (TrieNode n : tmp.edge) {
                        tmp = null;
                        if (n != null) {
                            tmp = n;
                            break;
                        }
                    }
                }
                return;
            }
            current = tmp;
        }
        current.end--;
    }

    /** 获取word 是否存在 */
    public boolean wordExists(String word) {
        TrieNode tmp = null;
        return (tmp = getWordNode(word)) != null && tmp.end > 0;
    }

    /** 获取word 所在的节点 */
    public TrieNode getWordNode(String word) {
        // null 不处理
        if (word == null) {
            return null;
        }
        // 空字符串
        if (word.length() == 0) {
            return this.head;
        }
        // 有字符串
        TrieNode current = this.head;
        for (char c : word.toCharArray()) {
            TrieNode tmp = current.edge[getAscallLetterIndex(c)];
            if (tmp == null) {
                return null;
            }
            current = tmp;
        }
        return current;
    }

    /** 获取 word 出现的次数 */
    public int getWordNum(String word) {
        TrieNode node = getWordNode(word);
        return node == null ? 0 : node.end;
    }

    /** 获取以word为前缀的单词出现的次数 */
    public int getPrifixWordNum(String word) {
        TrieNode node = getWordNode(word);
        return node == null ? 0 : node.pass;
    }

    /** 单词总数 */
    public int size() {
        return head.pass;
    }

    /** 出现最多的单词节点是 */
    public TrieNode getMaxWord() {
        if (this.nodes.size() == 0) {
            return null;
        }
        this.nodes.sort(Comparator.comparingInt(n -> n.end));
        return this.nodes.get(this.nodes.size() - 1);
    }

    /** 出现最少的单词节点是 */
    public TrieNode getMinWord() {
        if (this.nodes.size() == 0) {
            return null;
        }
        this.nodes.sort(Comparator.comparingInt(n -> n.end));
        return this.nodes.get(0);
    }

    @SuppressWarnings("all")
    public static int getAscallLetterIndex(char c) {
        int t = c;
        if (t > 127) {
            throw new IllegalArgumentException("字符不被允许！");
        }
        return t;
    }

    @Getter
    @Setter
    public static class TrieNode {
        public String word;
        // 经过数
        public int pass;
        // 抵达结束数
        public int end;
        public TrieNode[] edge = new TrieNode[128];
    }
}
